Unconditionally pulling the image before running means this script will now fail in offline/air-gapped environments (or when using a locally-built image tag) even if the image is already present locally. Consider making the pull optional (e.g., behind a PULL_DOCKER_IMAGE flag) or using a runtime option like "--pull=missing" where supported.

CLEAN_DOCKER_CONTAINER=1 will cause examples/run_local_pretrain.sh to remove all containers on the host (it runs "docker/podman ps -aq" and rm -f each). That’s a risky default for shared nodes; consider defaulting this to 0 and only enabling it explicitly when you’re sure it’s safe.

PRIMUS_EXP_NAME is assigned twice; the second assignment overrides the descriptive name built from the run parameters, so logs/output will always go to the hard-coded "debugv3" directory. Remove the second export or gate it behind an explicit debug flag to avoid accidental output collisions.

The header comment says the script "downloads" C4 shards and that the default is 1 shard, but the implementation explicitly skips downloading and defaults NUM_SHARDS to 200. Please align the documentation and defaults with the actual behavior (either implement download, or update comments and set NUM_SHARDS default accordingly).

This merge step appends directly into the final JSONL_FILE with '>>'. If the script is interrupted or a shard is corrupt, you'll end up with a partial JSONL that then causes future runs to skip merging because the file exists. Write to a temporary file and atomically move it into place on success (and/or validate the output) to avoid leaving a bad cached artifact.

CLEAN_DOCKER_CONTAINER=1 will cause examples/run_local_pretrain.sh to remove all containers on the host (it runs "docker/podman ps -aq" and rm -f each). That’s a risky default for shared nodes; consider defaulting this to 0 and only enabling it explicitly when you’re sure it’s safe.

usage_tensor is created as float32 even though r_used is a byte count on the order of 1e11 for large GPUs. float32 will lose integer precision at that scale, which can make max-rank selection and reported GB values slightly inaccurate. Prefer using an integer dtype (e.g., int64) for the gathered byte counts, then format as GB on CPU.

This script exports placeholder values for HF_TOKEN and WANDB_API_KEY, which will overwrite any real credentials already set in the environment and can cause training/CI runs to fail unless the file is edited. Prefer reading from the environment and erroring if unset (or only setting defaults when variables are empty), rather than hard-coding placeholder strings.

PRIMUS_PATH defaults to a user-specific absolute path (/shared/john/Primus), which makes this script non-portable and likely to fail for other users/environments. Consider requiring PRIMUS_PATH to be provided (and exit with a clear message if unset) or deriving it relative to the repo root.

The unit test jobs for both torch and jax are commented out, which removes automated test execution from CI and can allow regressions (including in this PR) to merge unnoticed. If this is temporary, consider gating them behind a workflow input/label, reducing the test scope, or moving them to a separate workflow instead of disabling them entirely.

Passing --nodelist with an empty value (when SLURM_NODELIST is unset) causes srun to fail with an invalid nodelist. Consider only adding --nodelist when SLURM_NODELIST is non-empty (or default to omitting the flag entirely).

This hard-codes an extra bind mount to /shared_aig/c4, which will break (or unintentionally create an empty host directory) on systems that don’t have that path. Consider making this conditional/configurable via an env var (e.g., EXTRA_VOLUME_ARGS / PRIMUS_EXTRA_MOUNTS) rather than always mounting a cluster-specific path.

The header comment says the default downloads 1 shard for testing, but NUM_SHARDS is set to default to 200. Please update either the comment or the default so they match (and avoid surprising users by downloading/processing 200 shards by default).

subprocess.run(..., check=True) will raise an exception on failure, so the subsequent if ret.returncode != 0: branch is redundant/unreachable. Either remove the returncode check, or set check=False and keep the explicit returncode handling (and ideally invoke pip via sys.executable -m pip to ensure the active Python environment is used).

The Torch/JAX unit test jobs are entirely commented out, which effectively disables CI test execution for PRs and can allow regressions to merge undetected. If this is only needed temporarily, consider gating the jobs behind an if: condition (e.g., on a label/branch) rather than removing them from the workflow, or re-enable them before merging.

This script exports HF_TOKEN and WANDB_API_KEY as literal strings, which both encourages storing secrets in a committed file and also overwrites any values already present in the caller’s environment. Prefer reading these from the environment (and erroring if missing) rather than exporting placeholder values here.

max_usage / r_total uses the local rank’s r_total when computing the max-usage ratio, which can be incorrect if ranks have different total HBM sizes (or if totals differ for any reason). To make this accurate, gather each rank’s r_total (or gather precomputed r_used/r_total ratios) and compute the max ratio corresponding to max_rank.

This config changes the default pipeline_model_parallel_size from the common default of 1 to 8, which is inconsistent with the MI355X DeepSeek-V3 FP8 config and most other configs. To avoid surprising users (and breaking single-node runs), consider keeping the default at 1 and relying on PRIMUS_PP to opt into PP>1.